# 引入scipy中的距离函数，默认欧式距离
import sklearn.metrics
from matplotlib import pyplot as plt
from scipy.spatial.distance import cdist
import numpy as np


class KMeans:
    # 初始化，参数n_clusters(指定簇的个数)、max_iter（最大迭代次数）
    def __init__(self, n_clusters=6, max_iter=300):
        self.n_clusters = n_clusters
        self.max_iter = max_iter
        self.centroids = None

    # 训练模型方法，K-Means聚类过程，传入原始数据
    def fit(self, datas):
        # 随机选取datas中的点作为初始的簇中心点
        n = datas.shape[0]  # 数据集datas的对象（样本点）的总数量
        indexes_centroid = np.random.choice(np.arange(0, n), size=self.n_clusters, replace=False)
        self.centroids = datas[indexes_centroid, :]  # 选取随机抽取的n_clusters个簇中心点

        # 开始迭代
        c_index = np.zeros((n,))
        for i in range(self.max_iter):
            # 1.计算距离矩阵，得到的是一个 n * n_clusters 的矩阵
            # 每一行代表一个样本点到所有中心的距离，一行里n_clusters个值分别指到第几个中心的距离
            distances = cdist(datas, self.centroids)  # 两个数组对象使用cdist计算欧式距离，要求它们维度一致

            # 2.选取最近的中心点的类别作为当前点的分类
            c_index = np.argmin(distances, axis=1)  # axis=1每一行取最小值，最后结果保存为一列（ n*1 的矩阵）

            # 3.对当前每一类数据进行均值计算，更新簇的中心点
            for i in range(self.n_clusters):  # i为 0 ~ n_clusters-1
                # 首先排除掉没有出现在c_index里的类别（即所有的点都没有离这个中心最近）
                if i not in c_index:
                    continue   # 跳过是为了防止错误更新 无需更新的簇中心点

                # 选出所有类别是i的点，取datas里面坐标的均值，更新第i个中心
                # c_index==i逻辑判断表达式，结果为布尔类型（数组）
                # c_index为一个数组，datas[c_index==i]返回结果为true对应的data的值，即类别为i的值的坐标
                self.centroids[i] = np.mean(datas[c_index == i], axis=0)  # 布尔索引。axis=0得到一行的数据，将每一列做均值计算，列数不变

        return c_index  # 返回最后的数据集每个对象的簇中心索引

    # 实现预测方法
    def predict(self, samples):  # samples一组样本点（新来的测试数据）
        # 跟上面一样，先计算距离矩阵，然后选取距离最近的那个中心的类别
        distances = cdist(samples, self.centroids)
        c_index = np.argmin(distances, axis=1)
        return c_index


class KMeansPlusPlus(KMeans):
    # 参数n_clusters(指定簇的个数)、max_iter（最大迭代次数）、centroids（可以指定初始中心，若不指定将执行K-Means++的算法获得）
    def __init__(self, n_clusters=6, max_iter=300, centroids: list = None):
        super().__init__(n_clusters, max_iter)
        self.centroids = np.array(centroids)

    def initialize_centroids(self, X):
        """
        使用 K-Means++ 算法初始化聚类中心
        :param X: 数据集，形状为 (n_samples, n_features)
        """
        # 选择第一个聚类中心作为随机样本点
        n_samples = X.shape[0]
        centroids = [X[np.random.randint(n_samples)]]

        # 选择剩余的 n_clusters - 1 个聚类中心
        for k in range(1, self.n_clusters):
            # 计算所有点到最近聚类中心的距离的平方(np.inner(c - x, c - x)计算向量差的点积，即计算了向量差的平方和，也即(c - x) · (c - x))
            distances = np.array([min([np.inner(c - x, c - x) for c in centroids]) for x in X])

            # 概率与距离的平方成反比，因此使用距离的平方的倒数来计算累积概率
            probabilities = distances / distances.sum()
            cumulative_probabilities = probabilities.cumsum()

            # 通过轮盘赌方法选择下一个聚类中心
            i = k
            r = np.random.rand()  # 从 [0.0, 1.0) 区间内生成均匀分布的随机数
            for j, p in enumerate(cumulative_probabilities):
                if r < p:  # “距离当前i - 1个簇中心点越远的点会有更高的概率被选中” 的体现
                    i = j
                    break

                    # 添加新的聚类中心
            centroids.append(X[i])

        self.centroids = np.array(centroids)

    # 训练模型方法，K-Means++聚类过程，传入原始数据
    def fit(self, datas):
        n = datas.shape[0]  # 数据集datas的对象（样本点）的总数量
        if self.centroids is None or len(self.centroids.shape) == 0:   # 执行K-Means++ 的选取初始中心点的方法
            self.initialize_centroids(datas)

        # 开始迭代
        c_index = np.zeros((n,))
        for i in range(self.max_iter):
            # 1.计算距离矩阵，得到的是一个 n * n_clusters 的矩阵
            # 每一行代表一个样本点到所有中心的距离，一行里n_clusters个值分别指到第几个中心的距离
            distances = cdist(datas, self.centroids)  # 两个数组对象使用cdist计算欧式距离，要求它们维度一致

            # 2.选取最近的中心点的类别作为当前点的分类
            c_index = np.argmin(distances, axis=1)  # axis=1每一行取最小值，最后结果保存为一列（ n*1 的矩阵）

            # 3.对当前每一类数据进行均值计算，更新簇的中心点
            for i in range(self.n_clusters):  # i为 0 ~ n_clusters-1
                # 首先排除掉没有出现在c_index里的类别（即所有的点都没有离这个中心最近）
                if i not in c_index:
                    continue  # 跳过是为了防止错误更新 无需更新的簇中心点

                # 选出所有类别是i的点，取datas里面坐标的均值，更新第i个中心
                # c_index==i逻辑判断表达式，结果为布尔类型（数组）
                # c_index为一个数组，datas[c_index==i]返回结果为true对应的data的值，即类别为i的值的坐标
                self.centroids[i] = np.mean(datas[c_index == i], axis=0)  # 布尔索引。axis=0得到一行的数据，将每一列做均值计算，列数不变

        return c_index  # 返回最后的数据集每个对象的簇中心索引


if __name__ == '__main__':
    from loadDatas import loadDatas
    from sklearn.metrics import adjusted_rand_score

    X_train, X_test, y_train, y_test = loadDatas('iris', 0.0001)

    kmeans1 = KMeans(n_clusters=3, max_iter=200)
    labels1 = kmeans1.fit(X_train)
    ARI = adjusted_rand_score(y_train, labels1)
    print('K-Means\n\t ARI = ', ARI)

    centroids = []
    for i in range(3):   # 给出三个不同类别的中心点作为簇的初始中心点
        for j, y in enumerate(y_train):
            if i == y:
                centroids.append(X_train[j])
                break
    kmeans2 = KMeansPlusPlus(n_clusters=3, max_iter=200, centroids=centroids)
    labels2 = kmeans2.fit(X_train)
    ARI = adjusted_rand_score(y_train, labels2)
    print('K-Means with centroids\n\t ARI = ', ARI)

    kmeans3 = KMeansPlusPlus(n_clusters=3, max_iter=200)
    labels3 = kmeans3.fit(X_train)
    ARI = adjusted_rand_score(y_train, labels3)
    print('K-Means++\n\tARI = ', ARI)